package com.chenjl.trace.client;

import java.util.HashMap;
import java.util.Map;

import org.slf4j.MDC;

import com.alibaba.dubbo.rpc.Filter;
import com.alibaba.dubbo.rpc.Invocation;
import com.alibaba.dubbo.rpc.Invoker;
import com.alibaba.dubbo.rpc.Result;
import com.alibaba.dubbo.rpc.RpcContext;
import com.alibaba.dubbo.rpc.RpcException;
import com.chenjl.trace.Tracer;
import com.chenjl.trace.enums.AnnotationEnum;
import com.chenjl.trace.enums.SpanTypeEnum;
import com.chenjl.trace.model.Endpoint;
import com.chenjl.trace.model.Span;
import com.chenjl.trace.utils.Constant;
/**
 * 跟踪调用记录filter，支持Dubbo RPC框架
 * 2018-6-29 15:34:24
 * @author chenjinlong
 */
public class DubboTraceFilter implements Filter {

	@Override
	public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {		
		RpcContext rpcContext = RpcContext.getContext();
		
		if(rpcContext.isConsumerSide()) {
			return this.consumerSideTrace(invoker,invocation,rpcContext);
		}
		else {
			return this.providerSideTrace(invoker,invocation,rpcContext);
		}
	}
	
	/**
	 * RPC消费者端trace
	 * @param invoker
	 * @param invocation
	 * @param rpcContext
	 * @return
	 */
	private final Result consumerSideTrace(Invoker<?> invoker, Invocation invocation,RpcContext rpcContext) {
		//dubbo服务消费者方
		Span rootSpan = Tracer.getRootSpan();
		if(rootSpan == null) {
			return invoker.invoke(invocation);
		}
		
		String serviceName = invoker.getInterface().getName();
		String methodName = invocation.getMethodName();
		
		Endpoint endpoint = new Endpoint();
		endpoint.setServiceName("dubbo service");
		endpoint.setHost(rpcContext.getRemoteHost());
		endpoint.setPort(rpcContext.getRemotePort());
		String spanId = Tracer.begin(rootSpan.getTraceId(),rootSpan.getId(),serviceName+"."+methodName,endpoint,AnnotationEnum.CLIENT_SEND,SpanTypeEnum.DUBBO_SERVICE);
		
		rpcContext.setAttachment(Constant.TRACE_ID, rootSpan.getTraceId());
		rpcContext.setAttachment(Constant.TRACE_PARENT_ID,spanId);
		
		Result invokeResult = invoker.invoke(invocation);
		
		boolean hasExceptionFlag = invokeResult.hasException();
		Map<String,String> datas = new HashMap<String,String>(2);
		datas.put("hasException",String.valueOf(hasExceptionFlag));
		Tracer.end(spanId,endpoint,AnnotationEnum.CLIENT_RECEIVE,datas);
		return invokeResult;
	}
	/**
	 * RPC提供者端trace
	 * @param invoker
	 * @param invocation
	 * @param rpcContext
	 * @return
	 */
	private final Result providerSideTrace(Invoker<?> invoker, Invocation invocation,RpcContext rpcContext) {
		//dubbo服务提供者方
		String traceId = rpcContext.getAttachment(Constant.TRACE_ID);
		String parentId = rpcContext.getAttachment(Constant.TRACE_PARENT_ID);
		if(traceId == null || parentId == null) {
			return invoker.invoke(invocation);
		}
		
		String serviceName = invoker.getInterface().getName();
		String methodName = invocation.getMethodName();
		
		MDC.put(Constant.C_TRACE_ID,traceId);
		
		Endpoint endpoint = new Endpoint();
		endpoint.setServiceName("dubbo service");
		endpoint.setHost(rpcContext.getRemoteHost());
		endpoint.setPort(rpcContext.getRemotePort());
		
		String spanId = Tracer.begin(traceId,parentId,serviceName+"."+methodName,endpoint,AnnotationEnum.SERVER_RECEIVE,SpanTypeEnum.DUBBO_SERVICE);
		
		Result invokeResult = invoker.invoke(invocation);
		
		boolean hasExceptionFlag = invokeResult.hasException();
		Map<String,String> datas = new HashMap<String,String>(2);
		datas.put("hasException",String.valueOf(hasExceptionFlag));
		
		Tracer.end(spanId,endpoint,AnnotationEnum.SERVER_SEND,datas);
		Tracer.clear();
		return invokeResult;
	}
}